#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <openssl/md5.h>

#define SERVER_IP "127.0.0.1"
#define SERVER_PORT 8080
#define BUFFER_SIZE 4096
#define MD5_DIGEST_LENGTH 16

// 连接到服务器
int connect_to_server() {
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("socket creation failed");
        return -1;
    }
    
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(SERVER_PORT);
    
    if (inet_pton(AF_INET, SERVER_IP, &server_addr.sin_addr) <= 0) {
        perror("invalid address");
        close(sock);
        return -1;
    }
    
    if (connect(sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("connection failed");
        close(sock);
        return -1;
    }
    
    return sock;
}

// 序列化字符串（长度前缀法）
int serialize_string(const char* str, uint8_t* buffer, size_t buffer_size) {
    size_t len = strlen(str);
    if (len + sizeof(uint32_t) > buffer_size) {
        return -1; // 缓冲区太小
    }
    
    // 先写入字符串长度（大端序）
    uint32_t net_len = htonl((uint32_t)len);
    memcpy(buffer, &net_len, sizeof(net_len));
    
    // 再写入字符串内容
    memcpy(buffer + sizeof(net_len), str, len);
    
    return sizeof(net_len) + len;
}

// 发送文件下载请求
int send_file_request(int sock, const char* filename) {
    uint8_t buffer[BUFFER_SIZE];
    
    // 序列化文件名
    int filename_len = serialize_string(filename, buffer, BUFFER_SIZE);
    if (filename_len < 0) {
        printf("文件名太长\n");
        return -1;
    }
    
    // 发送请求
    if (send(sock, buffer, filename_len, 0) < 0) {
        perror("发送请求失败");
        return -1;
    }
    
    return 0;
}

// 接收服务器响应
int receive_file_info(int sock, char* filename, size_t filename_size, 
                     size_t* file_size, uint8_t* md5_digest) {
    uint8_t buffer[BUFFER_SIZE];
    int total_received = 0;
    
    // 接收文件名长度
    while (total_received < sizeof(uint32_t)) {
        int received = recv(sock, buffer + total_received, 
                           sizeof(uint32_t) - total_received, 0);
        if (received <= 0) {
            perror("接收文件名长度失败");
            return -1;
        }
        total_received += received;
    }
    
    // 解析文件名长度
    uint32_t net_filename_len;
    memcpy(&net_filename_len, buffer, sizeof(net_filename_len));
    size_t filename_len = ntohl(net_filename_len);
    
    // 接收文件名
    total_received = 0;
    while (total_received < filename_len) {
        int received = recv(sock, buffer + sizeof(uint32_t) + total_received, 
                           filename_len - total_received, 0);
        if (received <= 0) {
            perror("接收文件名失败");
            return -1;
        }
        total_received += received;
    }
    
    // 确保文件名以NULL结尾
    size_t copy_len = filename_len < filename_size - 1 ? filename_len : filename_size - 1;
    memcpy(filename, buffer + sizeof(uint32_t), copy_len);
    filename[copy_len] = '\0';
    
    // 接收文件大小
    total_received = 0;
    while (total_received < sizeof(uint64_t)) {
        int received = recv(sock, (uint8_t*)file_size + total_received, 
                           sizeof(uint64_t) - total_received, 0);
        if (received <= 0) {
            perror("接收文件大小失败");
            return -1;
        }
        total_received += received;
    }
    
    // 转换文件大小字节序
    *file_size = be64toh(*file_size);
    
    // 接收MD5摘要
    total_received = 0;
    while (total_received < MD5_DIGEST_LENGTH) {
        int received = recv(sock, md5_digest + total_received, 
                           MD5_DIGEST_LENGTH - total_received, 0);
        if (received <= 0) {
            perror("接收MD5摘要失败");
            return -1;
        }
        total_received += received;
    }
    
    return 0;
}

// 接收文件内容
int receive_file_content(int sock, const char* filename, size_t file_size, 
                        const uint8_t* expected_md5) {
    FILE* file = fopen(filename, "wb");
    if (!file) {
        perror("创建文件失败");
        return -1;
    }
    
    uint8_t buffer[BUFFER_SIZE];
    size_t total_received = 0;
    MD5_CTX md5_ctx;
    uint8_t actual_md5[MD5_DIGEST_LENGTH];
    
    // 初始化MD5计算
    MD5_Init(&md5_ctx);
    
    printf("开始下载文件: %s (大小: %zu bytes)\n", filename, file_size);
    printf("进度: ");
    
    // 接收文件内容
    while (total_received < file_size) {
        size_t to_receive = file_size - total_received;
        if (to_receive > BUFFER_SIZE) {
            to_receive = BUFFER_SIZE;
        }
        
        int received = recv(sock, buffer, to_receive, 0);
        if (received <= 0) {
            perror("接收文件内容失败");
            fclose(file);
            return -1;
        }
        
        // 写入文件
        if (fwrite(buffer, 1, received, file) != received) {
            perror("写入文件失败");
            fclose(file);
            return -1;
        }
        
        // 更新MD5计算
        MD5_Update(&md5_ctx, buffer, received);
        
        total_received += received;
        
        // 显示进度
        int progress = (int)((double)total_received / file_size * 100);
        printf("\r进度: %d%%", progress);
        fflush(stdout);
    }
    
    printf("\n下载完成\n");
    
    // 完成MD5计算
    MD5_Final(actual_md5, &md5_ctx);
    
    // 验证MD5
    if (memcmp(expected_md5, actual_md5, MD5_DIGEST_LENGTH) == 0) {
        printf("MD5验证成功\n");
    } else {
        printf("MD5验证失败\n");
        printf("期望的MD5: ");
        for (int i = 0; i < MD5_DIGEST_LENGTH; i++) {
            printf("%02x", expected_md5[i]);
        }
        printf("\n实际的MD5: ");
        for (int i = 0; i < MD5_DIGEST_LENGTH; i++) {
            printf("%02x", actual_md5[i]);
        }
        printf("\n");
    }
    
    fclose(file);
    return 0;
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        printf("用法: %s <文件名>\n", argv[0]);
        return 1;
    }
    
    const char* filename = argv[1];
    
    // 连接到服务器
    int sock = connect_to_server();
    if (sock < 0) {
        return 1;
    }
    
    printf("已连接到服务器\n");
    
    // 发送文件下载请求
    if (send_file_request(sock, filename) < 0) {
        close(sock);
        return 1;
    }
    
    printf("已发送文件请求: %s\n", filename);
    
    // 接收文件信息
    char received_filename[256];
    size_t file_size;
    uint8_t md5_digest[MD5_DIGEST_LENGTH];
    
    if (receive_file_info(sock, received_filename, sizeof(received_filename), 
                         &file_size, md5_digest) < 0) {
        close(sock);
        return 1;
    }
    
    printf("文件信息接收成功:\n");
    printf("  文件名: %s\n", received_filename);
    printf("  文件大小: %zu bytes\n", file_size);
    printf("  MD5: ");
    for (int i = 0; i < MD5_DIGEST_LENGTH; i++) {
        printf("%02x", md5_digest[i]);
    }
    printf("\n");
    
    // 接收文件内容
    if (receive_file_content(sock, received_filename, file_size, md5_digest) < 0) {
        close(sock);
        return 1;
    }
    
    close(sock);
    return 0;
}